use ndarray::{prelude::*, ArcArray2};

pub trait AddFront {
    fn add_front(&self, axis: Axis, new: ArrayView1<'_, f64>) -> Array2<f64>;
}

impl AddFront for Array2<f64> {
    fn add_front(&self, axis: Axis, value: ArrayView1<'_, f64>) -> Array2<f64> {
        let rows = self.len_of(Axis(0));
        let cols = self.len_of(Axis(1));
        match axis {
            Axis(0) => {
                let mut new = Array2::zeros((0, cols));
                new.push_row(value).unwrap();
                for r in self.rows() {
                    new.push_row(r).unwrap();
                }
                new
            }

            Axis(1) => {
                let mut new = Array2::zeros((rows, 0));
                new.push_column(value).unwrap();
                for c in self.columns() {
                    new.push_column(c).unwrap();
                }
                new
            }
            _ => panic!("Axis must be 0 or 1"),
        }
    }
}

impl AddFront for ArrayView2<'_, f64> {
    fn add_front(&self, axis: Axis, value: ArrayView1<'_, f64>) -> Array2<f64> {
        let rows = self.len_of(Axis(0));
        let cols = self.len_of(Axis(1));
        match axis {
            Axis(0) => {
                let mut new = Array2::zeros((0, cols));
                new.push_row(value).unwrap();
                for r in self.rows() {
                    new.push_row(r).unwrap();
                }
                new
            }

            Axis(1) => {
                let mut new = Array2::zeros((rows, 0));
                new.push_column(value).unwrap();
                for c in self.columns() {
                    new.push_column(c).unwrap();
                }
                new
            }
            _ => panic!("Axis must be 0 or 1"),
        }
    }
}

impl AddFront for ArcArray2<f64> {
    fn add_front(&self, axis: Axis, value: ArrayView1<'_, f64>) -> Array2<f64> {
        let rows = self.len_of(Axis(0));
        let cols = self.len_of(Axis(1));
        match axis {
            Axis(0) => {
                let mut new = Array2::zeros((0, cols));
                new.push_row(value).unwrap();
                for r in self.rows() {
                    new.push_row(r).unwrap();
                }
                new
            }

            Axis(1) => {
                let mut new = Array2::zeros((rows, 0));
                new.push_column(value).unwrap();
                for c in self.columns() {
                    new.push_column(c).unwrap();
                }
                new
            }
            _ => panic!("Axis must be 0 or 1"),
        }
    }
}
